#!/usr/bin/env python
# coding:utf-8
from PoboAPI import *
import datetime
import numpy as np

# RSI指标交易策略示例
# 策略逻辑：当RSI指标低于30时买入，高于70时卖出

# 开始时间，用于初始化参数
def OnStart(context):
    print("策略启动...")
    # 登录交易账号
    context.myacc = None
    if "回测期货" in context.accounts:
        print("登录交易账号[回测期货]")
        if context.accounts["回测期货"].Login():
            context.myacc = context.accounts["回测期货"]
    
    # 设置交易品种为螺纹钢主力合约
    g.exchange = 'SHFE'
    g.product = 'rb'
    g.code = GetMainContract(g.exchange, g.product, 20)
    print(f"初始主力合约：{g.code}")
    
    # 设置RSI参数
    g.rsi_period = 14
    g.oversold = 30    # 超卖阈值
    g.overbought = 70  # 超买阈值
    
    # 设置持仓标志
    g.position = 0  # 0表示空仓，1表示多头，-1表示空头
    
    # 订阅K线数据
    SubscribeBar(g.code, BarType.Min30)

# 行情初始化事件
def OnMarketQuotationInitialEx(context, exchange, daynight):
    # 过滤不需要的交易所
    if exchange != g.exchange:
        return
    
    # 检查主力合约是否变更
    new_code = GetMainContract(g.exchange, g.product, 20)
    if g.code != new_code:
        print(f"主力合约变更：从 {g.code} 变为 {new_code}")
        # 取消旧合约订阅
        UnsubscribeBar(g.code, BarType.Min30)
        # 更新合约代码
        g.code = new_code
        # 订阅新合约
        SubscribeBar(g.code, BarType.Min30)
        
        # 如果有持仓，平掉旧合约的仓位
        if g.position != 0 and context.myacc is not None:
            price_type = PriceType(PbPriceType.Limit, 16, 0)
            if g.position > 0:
                QuickInsertOrder(context.myacc, g.code, 'sell', 'close', price_type, abs(g.position))
            else:
                QuickInsertOrder(context.myacc, g.code, 'buy', 'close', price_type, abs(g.position))
            g.position = 0

# K线事件
def OnBar(context, code, bartype):
    # 过滤不需要的合约
    if code != g.code:
        return
    
    # 获取历史数据
    df = GetHisDataAsDF(code, bar_type=BarType.Min30)
    if len(df) < g.rsi_period + 1:
        print("历史数据不足，等待更多数据...")
        return
    
    # 计算RSI指标
    close_prices = df['close'].values
    delta = np.diff(close_prices)
    gain = np.where(delta > 0, delta, 0)
    loss = np.where(delta < 0, -delta, 0)
    
    avg_gain = np.mean(gain[:g.rsi_period])
    avg_loss = np.mean(loss[:g.rsi_period])
    
    for i in range(g.rsi_period, len(delta)):
        avg_gain = (avg_gain * (g.rsi_period - 1) + gain[i]) / g.rsi_period
        avg_loss = (avg_loss * (g.rsi_period - 1) + loss[i]) / g.rsi_period
    
    if avg_loss == 0:
        rsi = 100
    else:
        rs = avg_gain / avg_loss
        rsi = 100 - (100 / (1 + rs))
    
    print(f"当前RSI值: {rsi:.2f}")
    
    # 交易逻辑
    if context.myacc is None:
        print("交易账户未登录，无法交易")
        return
    
    price_type = PriceType(PbPriceType.Limit, 16, 0)
    current_price = GetQuote(g.code).now
    
    # RSI低于超卖线且当前无持仓，买入开仓
    if rsi < g.oversold and g.position == 0:
        print(f"RSI低于{g.oversold}，买入开仓")
        volume = 1  # 交易手数，实际应用中可根据资金比例计算
        QuickInsertOrder(context.myacc, g.code, 'buy', 'open', price_type, volume)
        g.position = volume
    
    # RSI高于超买线且当前为多头持仓，卖出平仓
    elif rsi > g.overbought and g.position > 0:
        print(f"RSI高于{g.overbought}，卖出平仓")
        QuickInsertOrder(context.myacc, g.code, 'sell', 'close', price_type, g.position)
        g.position = 0
    
    # 每日收盘前平仓逻辑可以通过OnAlarm事件实现
    # 这里省略了该部分代码

# 设置收盘前平仓的闹钟事件
def OnAlarm(context, alarmid):
    # 检查是否有持仓
    if g.position != 0 and context.myacc is not None:
        print("收盘前平仓")
        price_type = PriceType(PbPriceType.Limit, 16, 0)
        if g.position > 0:
            QuickInsertOrder(context.myacc, g.code, 'sell', 'close', price_type, g.position)
        else:
            QuickInsertOrder(context.myacc, g.code, 'buy', 'close', price_type, abs(g.position))
        g.position = 0